import os
from CalTorque import CalTorque
from datetime import datetime
import matplotlib
from matplotlib import pyplot as PLT
from scipy.stats import linregress
import numpy

data_files = []
dp = 'data/'
ls= os.listdir(dp)
for i in ls:
   if i.split('.')[-1] not in ['xlsx', 'xls']: continue
   if 'out' in i: continue
   elif 'in' not in i: continue
   data_files.append(dp+i)

bins = 400
max_d = 400

max_torque = []
max_vel = []
max_at_B2 = []
max_at_B3 = []
filelist = []
avgtor = []
avgvel = []
linfit = []
linfit_B2 = []
linfit_B3 = []
linfit2 = []
linfit3 = []

for f in data_files:   
    data = CalTorque(f)
    if data['source'] != 'weakTh-228': continue
    if data['station'] != 'S5': continue
    if max(data['pos_offsetfix']) < 350: continue   # because some stations in run_history are wrong
    if max(data['pos_offsetfix']) > 400: continue
    print f
    
    #print len(data['rpm']), len(data['velocity'])
    #for i, item in enumerate(data['rpmvel']):
    #    print data['RPM'][i], '         ', data['rpm'][i], data['velocity'][i]
    
    filelist.append(f.split('.')[0].split(' ')[2].split('_')[0])
    max_torque.append(data['max_torque_inoz'])
    max_vel.append(max(data['velocity']))
    
    X = numpy.linspace(0,max_d,bins)
    T = numpy.zeros(bins)
    V = numpy.zeros(bins)
    
    # bin
    for j,x in enumerate(X):
        pos_torque = []
        pos_vel = []
        x0 = x-(max_d/bins)/2.
        x1 = x+(max_d/bins)/2.
        for i,p in enumerate(data['position']):
            if p>x0 and p<=x1:
                pos_torque.append(data['torque_inoz'][i])
        if len(pos_torque)>0:
            T[j] += max(pos_torque)
    
    # average
    t = 0
    for torque in T:
        t += torque
    t2 = t / len(T)
    avgtor.append(t2)
    
    sumvel = 0
    for j, v in enumerate(data['vel']):
        sumvel += v
    averagevel = sumvel / int(len(data['vel']))
    avgvel.append(averagevel)
   
    pos_torque = []
    for i, x in enumerate(data['pos_offsetfix']):
        if x > 300 and x < 330:
            pos_torque.append(data['torque_inoz'][i])
    max_at_B2.append(max(pos_torque))
    
    pos_torque2 = []
    for i, x in enumerate(data['pos_offsetfix']):
        if x > 350 and x < 375:
            pos_torque2.append(data['torque_inoz'][i])
    max_at_B3.append(max(pos_torque2))
    

lin = linregress(avgvel, max_torque)
lin_B2 = linregress(avgvel, max_at_B2)
lin_B3 = linregress(avgvel, max_at_B3)
lin2 = linregress(avgvel, avgtor)
lin3 = linregress(max_vel, max_torque)


def Polyfnc(x, m, b):
    return m*x + b

for thing in avgvel:
    linfit.append(Polyfnc(thing, lin[0], lin[1]))
    linfit_B2.append(Polyfnc(thing, lin_B2[0], lin_B2[1]))
    linfit_B3.append(Polyfnc(thing, lin_B3[0], lin_B3[1]))
    linfit2.append(Polyfnc(thing, lin2[0], lin2[1]))
for other in max_vel:
    linfit3.append(Polyfnc(other, lin3[0], lin3[1]))

box = [0.14, 0.14, 0.76, 0.76]

fig = PLT.figure(figsize=(15,8),dpi=150)
ax1 = fig.add_axes(box)
ax1.set_ylabel('Max Torque (in.oz)')
ax1.set_xlabel('Average Velocity (cm/s)')
PLT.title('Max Torque Versus Average Velocity')
ax1.scatter(avgvel, max_torque, color = 'g')
ax1.plot(avgvel, linfit, '-', color = 'g')
#for i, txt in enumerate(filelist):
    #ax1.annotate(txt, (date[i], max_torque[i]), position = (date[i], max_torque[i]), size = 'x-small', alpha = 0.5, rotation = 45)
fig.savefig('vel/maxt_vs_vel.png')

fig2 = PLT.figure(figsize=(15,8),dpi=150)
ax2 = fig2.add_axes(box)
ax2.set_ylabel('Max Torque (in.oz)')
ax2.set_xlabel('Date')
PLT.title('Max Torque Versus Average Velocity for Bend A6 and A7')
ax2.scatter(avgvel, max_at_B2, color = 'b', label = 'Bend A6')
ax2.plot(avgvel, linfit_B2, '-', color = 'b')
ax2.scatter(avgvel, max_at_B3, color = 'r', marker = 'D', label = 'Bend A7')
ax2.plot(avgvel, linfit_B3, '-', color = 'r')
fig2.savefig('vel/maxt_vs_vel_for_A67.png')

fig4 = PLT.figure(figsize=(15,8),dpi=150)
ax4 = fig4.add_axes(box)
ax4.set_ylabel('Average Torque (in.oz)')
ax4.set_xlabel('Average Velocity (cm/s)')
PLT.title('Average Torque Versus Average Velocity')
ax4.scatter(avgvel, avgtor, color = 'b')
ax4.plot(avgvel, linfit2, '-', color = 'b')
fig4.savefig('vel/avgt_vs_vel.png')

fig5 = PLT.figure(figsize=(15,8),dpi=150)
ax5 = fig5.add_axes(box)
ax5.set_ylabel('Maximum Torque (in.oz)')
ax5.set_xlabel('Maximum Velocity (cm/s)')
PLT.title('Maximum Torque Versus Maximum Velocity')
ax5.scatter(max_vel, max_torque, color = 'g')
fig5.savefig('vel/maxt_vs_maxv.png')
